import pandas as pd

stocks_path = r'd:\Python\study_data\stocks\\'

class Position():
    def __init__(self, stock_name, date, cnt):
        self.stock_name = stock_name
        self.buy_price = 0
        self.buy_date = date
        self.sell_price = 0
        self.sell_date = ''
        self.hold_days = 0
        self.profit = 0
        self.profit_rate = 0
        self.hold_cnt = 0

def backtest(all_buys, hold_days, hold_max):
    date_df = pd.read_csv(r'd:\Python\study_data\stocks\sz000001.csv')['date']
    all_buys_df = pd.read_csv(all_buys, index_col=0, low_memory=False).sort_index()
    position = {}
    cash = 1000000
    for i in range(len(date_df)):
        date = date_df.iloc[i]
        if date in all_buys_df.index:
            stocks = all_buys_df.loc[date].dropna().values
            buy_stocks_num = (len(stocks) + 1) // 2
            buy_stocks = stocks[:buy_stocks_num]
            for stock in buy_stocks:
                if stock not in position.keys():
                    if len(position) < hold_max:
                        position[stock] = Position(stock, date, i)
                        stock_df = pd.read_csv(stocks_path + stock + '.csv', index_col=0)
                        position[stock].buy_price = stock_df.loc[date, 'close']
                        position[stock].sell_price = stock_df.loc[date_df.iloc[i + hold_days], 'close']
                        position[stock].profit_rate = position[stock].sell_price / position[stock].buy_price - 1
                else:
                    position[stock].hold_cnt += 1
                    if position[stock].hold_cnt == hold_days:
                        cash = cash + cash / hold_max * position[stock].profit_rate
                        del position[stock]
        else:
            for stock in list(position.keys()):
                position[stock].hold_cnt += 1
                if position[stock].hold_cnt == hold_days:
                    cash = cash + cash / hold_max * position[stock].profit_rate
                    del position[stock]

        print(date, cash)

if __name__ == '__main__':
    backtest('all_buys_df.csv', 20, 20)